import io
import pickle
from flask import Blueprint, jsonify , request, send_file,render_template
from flask_restful import Resource
import pandas as pd
from models import Model,db,Data,User
from utils import get_model,pretreatment,train,metrics_calculate
from flask_jwt import jwt_required, current_identity
import json
from loguru import logger

model_bp = Blueprint("model", __name__, url_prefix="/model")

class ModelResource(Resource):
    @jwt_required()
    def get(self, id=None):
        if id:
            model = Model.query.get_or_404(id)
            result = {'id': model.id, 'name': model.name,"model":model.model,'package':model.package,"group":model.group}
        else:
            group = request.form.get("group")
            if group is None:
                models = Model.query.all()
            else:
                models = Model.query.filter(group==group).all()
            result = [{'id': model.id, 'name': model.name,"model":model.model,'package':model.package,"group":model.group} for model in models]
        return {"code":"200","msg":result}

    @jwt_required()
    def post(self):
        name = request.form.get("name")
        model = request.form.get("model")
        package = request.form.get("package")
        group = request.form.get("group")
        param = {}
        try:
            rparam = json.loads(request.form.get("param"))
            for i in rparam:
                param[i] = rparam[i]
        except:
            pass
        file = request.files.get('file')
        if file:
                data = file.read()
                data = pickle.loads(data)
        else:
            try:
                data = get_model(model,package,**rparam)
                print(data)
            except:
                return {"code":"201","msg":"生产模型错误"}
        new_model = Model(name=name,model=model,package=package,data=data,group=group )
        db.session.add(new_model)
        db.session.commit()
        return {"code":"200","msg":new_model.id}

    @jwt_required()
    def put(self, id):
        model = Model.query.get_or_404(id)
        model.name = request.form.get("name")
        model.model = request.form.get("model")
        model.package = request.form.get("package")
        file = request.files.get('file')
        if file:
            model.data = file.read()
        db.session.commit()
        return {"code":"200","msg":"修改成功"}

    @jwt_required()
    def delete(self, id):
        model = Model.query.get_or_404(id)
        db.session.delete(model)
        db.session.commit()
        return {"code":"200","msg":"删除成功"}

@jwt_required()
@model_bp.route('/download/<int:model_id>', methods=['GET'])
def download_data(model_id):
    model = Model.query.filter_by(id=model_id).first()
    if not model:
        return 'Model not found', 404

    filename = model.name + '.model'
    data = model.data
    return send_file(
        io.BytesIO(data),
        as_attachment=True,
        download_name=filename,
        mimetype='application/octet-stream'
    )

@jwt_required()
@model_bp.route('/train/<int:model_id>',methods=['POST'])
def train_model(model_id):
    model = Model.query.filter_by(id=model_id).first()
    if not model:
        return 'Model not found', 404
    m = model.get_data()
    data_id = request.form.get("dataId")
    data = Data.query.filter_by(id=data_id).first()
    df = data.get_data()
    traind,testd = pretreatment(df)
    logger.success(f'{model_id}开始训练')
    train(m,*traind)
    logger.success(f'{model_id}训练完成')
    model.data = pickle.dumps(m)
    db.session.commit()
    m = model.get_data()
    return metrics_calculate(m,*testd)

@jwt_required()
@model_bp.route('/predict/<int:model_id>',methods=['POST'])
def predict_model(model_id):
    model = Model.query.filter_by(id=model_id).first()
    if not model:
        return 'Model not found', 404
    m = model.get_data()
    data_id = request.form.get("dataId")
    data = Data.query.filter_by(id=data_id).first()
    if not data:
        return 'Data not found', 404
    df = data.get_data()
    X = pretreatment(df,label=False)
    Y = m.predict(X).tolist()
    r={}
    for i,y in zip(df['sample_id'],Y):
        r[i] = y
    return r

data_bp = Blueprint("data", __name__, url_prefix="/data")

class DataResource(Resource):
    
    @jwt_required()
    def get(self, id=None):
        if id:
            d = Data.query.get_or_404(id)
            df = d.get_data()
            info = {
                'class_name': str(df.__class__.__name__),
                'memory_usage': str(df.memory_usage(index=True, deep=True).sum()),
                'num_columns': str(len(df.columns)),
                'num_rows': str(len(df)),
                'total_size': str(df.values.nbytes),
                'cols': dict(zip(df.columns, df.dtypes.astype(str))),
            }
            result = {'id': d.id, 'name': d.name,"group":d.group,"info":info}
        else:
            group = request.form.get("group")
            if group is None:
                ds = Data.query.all()
            else:
                ds = Data.query.filter(group==group).all()
            result = [{'id': d.id, 'name': d.name,"group":d.group} for d in ds]
        return {"code":"200","msg":result}

    @jwt_required()
    def post(self):
        name = request.form.get("name")
        group = request.form.get("group")
        file = request.files.get('file')
        if not file:
            return {"code":"201","msg":"没有发现文件"}
        data = pd.read_csv(file)
        new_d = Data(name=name,data=data,group=group)
        db.session.add(new_d)
        db.session.commit()
        return {"code":"200","msg":new_d.id}

    @jwt_required()
    def put(self, id):
        d = Data.query.get_or_404(id)
        d.name = request.form.get("name")
        file = request.files.get('file')
        if file:
            d.data =  pickle.dumps(pd.read_csv(file))
        db.session.commit()
        return {"code":"200","msg":"修改成功"}

    @jwt_required()
    def delete(self, id):
        d = Data.query.get_or_404(id)
        db.session.delete(d)
        db.session.commit()
        return {"code":"200","msg":"删除成功"}


user_bp = Blueprint("user", __name__, url_prefix="/user")

class UserResource(Resource):
    
    @jwt_required()
    def get(self, id=None):
        if id:
            d = User.query.get_or_404(id)
            result = {'id': d.id, 'username': d.username,"password":d.password}
        else:
            ds = User.query.all()
            result = [{'id': d.id, 'username': d.username,"password":d.password} for d in ds]
        return {"code":"200","msg":result}

    def post(self):
        username = request.form.get("username")
        password = request.form.get("password")
        new_d = User(username=username,password=password)
        db.session.add(new_d)
        db.session.commit()
        return {"code":"200","msg":new_d.id}

    @jwt_required()
    def put(self, id):
        d = Data.query.get_or_404(id)
        d.password = request.form.get("password")
        db.session.commit()
        return {"code":"200","msg":"修改成功"}

    @jwt_required()
    def delete(self, id):
        d = Data.query.get_or_404(id)
        db.session.delete(d)
        db.session.commit()
        return {"code":"200","msg":"删除成功"}